from flask import Flask, request, redirect, url_for, render_template
import paddle
from PIL import Image
import numpy as np
import os

from werkzeug.utils import secure_filename

# 创建Flask应用实例
app = Flask(__name__)

# 设置上传文件的保存路径
UPLOAD_FOLDER = 'static'
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# 加载PaddlePaddle模型
model = paddle.vision.models.resnet18(pretrained=False)
model.fc = paddle.nn.Sequential(paddle.nn.Linear(512, 4), paddle.nn.Softmax(axis=1))
model_state_dict = paddle.load("work/Resnet18.pdparams")
model.set_state_dict(model_state_dict)
model.eval()

# 定义预处理函数
def preprocess_image(image_path):
    transform = paddle.vision.transforms.Compose([
        paddle.vision.transforms.Resize((224, 224)),
        paddle.vision.transforms.ToTensor(),
        paddle.vision.transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    image = Image.open(image_path)
    image = transform(image).unsqueeze(0)  # 增加一个批次维度
    return image

rice_diseases = ['Brown_Spot褐斑', 'Healthy健康', 'Leaf_Blast叶瘟', 'Neck_Blast茎瘟']

# 主页路由，渲染index.html模板文件
@app.route('/')
def index():
    return render_template('index.html')

# 处理图片预览路由，接收图片URL并返回预测结果
@app.route('/pre/<filename>')
def pre(filename):
    image_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    image = preprocess_image(image_path)
    predicts = model(image)
    predict_label = paddle.argmax(predicts, axis=1).numpy()[0]
    result = rice_diseases[predict_label]
    return render_template("result.html", result=result, image_url=filename)

# 处理文件上传路由
@app.route('/upload', methods=['POST'])
def upload_file():
    if request.method == 'POST':
        f = request.files['file']
        filename = secure_filename(f.filename)
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], filename)
        f.save(filepath)
        return redirect(url_for('pre', filename=filename))
    else:
        return render_template('index.html')

# 应用的主入口，启动Flask应用
if __name__ == '__main__':
    app.run(debug=True)
